# 创建时间：2024/9/3下午1:48
import csv
from enum import Enum
from prettytable import PrettyTable

persona = []
comp2 = []
comp3 = []


def init():
    """加载csv文件"""
    with open('./csv/persona.csv', 'r', encoding='utf-8') as cf:
        reader = csv.DictReader(cf)
        for row in reader:
            position = []
            if row['position1']:
                position.append(row['position1'])
            if row['position2']:
                position.append(row['position2'])
            if row['position3']:
                position.append(row['position3'])
            p = Persona(int(row['id']), row['arcana'], row['name'], int(row['level']), int(row['special']),
                        position)
            persona.append(p)
    with open('./csv/comp2.csv', 'r', encoding='utf-8') as cf:
        reader = csv.DictReader(cf)
        for row in reader:
            row['arcana1'] = Arcana.parse(row['arcana1'])
            row['arcana2'] = Arcana.parse(row['arcana2'])
            row['result'] = Arcana.parse(row['result'])
            comp2.append(row)
    with open('./csv/comp3.csv', 'r', encoding='utf-8') as cf:
        reader = csv.DictReader(cf)
        for row in reader:
            row['arcana1'] = Arcana.parse(row['arcana1'])
            row['arcana2'] = Arcana.parse(row['arcana2'])
            row['result'] = Arcana.parse(row['result'])
            comp3.append(row)


class Persona:
    def __init__(self, pid, arcana, name, level, special, position=None):
        self.pid = pid
        self.arcana = Arcana.parse(arcana)
        self.name = name
        self.level = level
        self.special = special
        self.position = position

    def __str__(self):
        return self.name

    def __repr__(self):
        return self.__str__()

    @staticmethod
    def parse_by_name(name):
        for p in persona:
            if p.name == name:
                return p
        else:
            return None


class PersonaNotFoundError(Exception):
    def __init__(self, name, message=None):
        if not message:
            message = f'未找到名称为"{name}"的人格面具，请确认名称'
        super().__init__(message)


class Arcana(Enum):
    FOOL = 0
    """愚者"""
    MAGICIAN = 1
    """魔术师"""
    PRIESTESS = 2
    """女教皇"""
    EMPRESS = 3
    """女皇"""
    EMPEROR = 4
    """皇帝"""
    HIEROPHANT = 5
    """教皇"""
    LOVERS = 6
    """恋爱"""
    CHARIOT = 7
    """战车"""
    JUSTICE = 8
    """正义"""
    HERMIT = 9
    """隐者"""
    FORTUNE = 10
    """命运"""
    STRENGTH = 11
    """力量"""
    HANGED = 12
    """倒悬者"""
    DEATH = 13
    """死神"""
    TEMPERANCE = 14
    """节制"""
    DEVIL = 15
    """恶魔"""
    TOWER = 16
    """塔"""
    STAR = 17
    """星"""
    MOON = 18
    """月亮"""
    SUN = 19
    """太阳"""
    JUDGMENT = 20
    """审判"""
    JOKER = 21
    """小丑"""
    ETERNALLY = 22
    """永劫"""
    WORLD = 23
    """世界"""

    def __str__(self):
        if self.name == 'FOOL':
            return '愚者'
        elif self.name == 'MAGICIAN':
            return '魔术师'
        elif self.name == 'PRIESTESS':
            return '女教皇'
        elif self.name == 'EMPRESS':
            return '女皇'
        elif self.name == 'EMPEROR':
            return '皇帝'
        elif self.name == 'HIEROPHANT':
            return '教皇'
        elif self.name == 'LOVERS':
            return '恋爱'
        elif self.name == 'CHARIOT':
            return '战车'
        elif self.name == 'JUSTICE':
            return '正义'
        elif self.name == 'HERMIT':
            return '隐者'
        elif self.name == 'FORTUNE':
            return '命运'
        elif self.name == 'STRENGTH':
            return '力量'
        elif self.name == 'HANGED':
            return '倒悬者'
        elif self.name == 'DEATH':
            return '死神'
        elif self.name == 'TEMPERANCE':
            return '节制'
        elif self.name == 'DEVIL':
            return '恶魔'
        elif self.name == 'TOWER':
            return '塔'
        elif self.name == 'STAR':
            return '星'
        elif self.name == 'MOON':
            return '月亮'
        elif self.name == 'SUN':
            return '太阳'
        elif self.name == 'JUDGMENT':
            return '审判'
        elif self.name == 'JOKER':
            return '小丑'
        elif self.name == 'ETERNALLY':
            return '永劫'
        elif self.name == 'WORLD':
            return '世界'

    def __repr__(self):
        return self.__str__()

    def __lt__(self, other):
        return self.value < other.value

    def __le__(self, other):
        return self.value <= other.value

    @classmethod
    def parse(cls, arcana):
        if arcana == '愚者':
            return cls.FOOL
        elif arcana == '魔术师':
            return cls.MAGICIAN
        elif arcana == '女教皇':
            return cls.PRIESTESS
        elif arcana == '女皇':
            return cls.EMPRESS
        elif arcana == '皇帝':
            return cls.EMPEROR
        elif arcana == '教皇':
            return cls.HIEROPHANT
        elif arcana == '恋爱':
            return cls.LOVERS
        elif arcana == '战车':
            return cls.CHARIOT
        elif arcana == '正义':
            return cls.JUSTICE
        elif arcana == '隐者':
            return cls.HERMIT
        elif arcana == '命运':
            return cls.FORTUNE
        elif arcana == '力量':
            return cls.STRENGTH
        elif arcana == '倒悬者':
            return cls.HANGED
        elif arcana == '死神':
            return cls.DEATH
        elif arcana == '节制':
            return cls.TEMPERANCE
        elif arcana == '恶魔':
            return cls.DEVIL
        elif arcana == '塔':
            return cls.TOWER
        elif arcana == '星':
            return cls.STAR
        elif arcana == '月亮':
            return cls.MOON
        elif arcana == '太阳':
            return cls.SUN
        elif arcana == '审判':
            return cls.JUDGMENT
        elif arcana == '小丑':
            return cls.JOKER
        elif arcana == '永劫':
            return cls.ETERNALLY
        elif arcana == '世界':
            return cls.WORLD
        elif arcana == '':
            return ''
        else:
            raise TypeError('输入有误，该阿尔卡那未定义')


class Calc:
    def __init__(self, target=None, material1=None, material2=None, material3=None):
        self._target = None
        self._material1 = None
        self._material2 = None
        self._material3 = None
        self._result = None
        self.target = target
        self.material1 = material1
        self.material2 = material2
        self.material3 = material3

    @property
    def target(self):
        return self._target

    @target.setter
    def target(self, target):
        if target:
            p = Persona.parse_by_name(target)
            if p:
                if p.special == 0:
                    self._target = p
                else:
                    raise TypeError('特殊合成的人格面具不能作为二体/三体合成的目标')
            else:
                raise PersonaNotFoundError(target)

    @property
    def material1(self):
        return self._material1

    @material1.setter
    def material1(self, m1):
        if m1:
            p = Persona.parse_by_name(m1)
            if p:
                self._material1 = p
            else:
                raise PersonaNotFoundError(m1)

    @material1.deleter
    def material1(self):
        self._material1 = None

    @property
    def material2(self):
        return self._material2

    @material2.setter
    def material2(self, m2):
        if m2:
            p = Persona.parse_by_name(m2)
            if p:
                self._material2 = p
            else:
                raise PersonaNotFoundError(m2)

    @material2.deleter
    def material2(self):
        self._material2 = None

    @property
    def material3(self):
        return self._material3

    @material3.setter
    def material3(self, m3):
        if m3:
            p = Persona.parse_by_name(m3)
            if p:
                self._material3 = p
            else:
                raise PersonaNotFoundError(m3)

    @material3.deleter
    def material3(self):
        self._material3 = None

    @property
    def result(self):
        return self._result

    @result.setter
    def result(self, result):
        self._result = result

    def show_result(self):
        result = self.result
        if not result:
            input('结果为空(回车确认)')
            return
        if len(result[0]) == 3:
            table = PrettyTable()
            table.field_names = ['阿尔卡那1', '人格面具1', '初始等级1', '阿尔卡那2', '人格面具2', '初始等级2']
            for row in result:
                p1 = row[1]
                p2 = row[2]
                table.add_row([p1.arcana, p1.name, p1.level, p2.arcana, p2.name, p2.level])
            print(table)
        elif len(result[0]) == 4:
            table = PrettyTable()
            table.field_names = ['阿尔卡那1', '人格面具1', '初始等级1', '阿尔卡那2', '人格面具2', '初始等级2',
                                 '阿尔卡那3', '人格面具3', '初始等级3']
            for row in result:
                p1 = row[1]
                p2 = row[2]
                p3 = row[3]
                table.add_row(
                    [p1.arcana, p1.name, p1.level, p2.arcana, p2.name, p2.level, p3.arcana, p3.name, p3.level])
            print(table)
        print(f'共计{len(result)}条结果', end='')
        input('(回车确认)')

    @staticmethod
    def showPersona(p):
        print('阿尔卡那:', p.arcana)
        print('人格面具:', p.name)
        print('初始等级:', str(p.level))
        print('掉落位置:', p.position)

    @staticmethod
    def _check_comp2(t, p1, p2):
        if t.pid == p1.pid or t.pid == p2.pid or p1.pid == p2.pid or t.special == 1:
            return False
        if p1.arcana != p2.arcana:
            # 不同arcana二体合成
            if p1.level + p2.level < t.level * 2:
                for p in persona:
                    if p.arcana != t.arcana:
                        continue
                    if p.pid >= t.pid:
                        return True
                    if p1.level + p2.level < p.level * 2:
                        return False
            return False
        else:
            # 同arcana二体合成
            if p1.level + p2.level >= t.level * 2:
                for p in persona:
                    if p.arcana != t.arcana or p.pid == t.pid or p.pid == p1.pid or p.pid == p2.pid or p.special == 1:
                        continue
                    if p.pid > t.pid and p1.level + p2.level >= p.level * 2:
                        return False
                    if p.pid > p1.pid and p.pid > p2.pid:
                        return True
            return False

    def calculate_persona_by_comp2_arithmetic(self):
        t = self.target
        m = self.material1
        result = []
        if t:
            # 通过合成目标计算素材
            if t.special == 1 or (m and m.pid == t.pid):
                self.result = result
                return
            list_p1 = []
            """可用的p1列表"""
            arcana_formula = []
            """可用的arcana公式列表"""
            # 计算可用的p1和arcana公式
            if m:
                # 如果指定了素材
                list_p1.append(m)
                for row in comp2:
                    if row['result'] == t.arcana and row['arcana1'] == m.arcana:
                        arcana_formula.append(row)
            else:
                # 如果未指定素材
                p1_arcana = set()
                for row in comp2:
                    if row['result'] == t.arcana and row['arcana1'] <= row['arcana2']:
                        arcana_formula.append(row)
                        p1_arcana.add(row['arcana1'])
                for p1 in persona:
                    if p1.arcana not in p1_arcana:
                        continue
                    list_p1.append(p1)
            for p1 in list_p1:
                p2_arcana = set()
                for row in arcana_formula:
                    if row['arcana1'] == p1.arcana:
                        p2_arcana.add(row['arcana2'])
                for p2 in persona:
                    if p2.arcana not in p2_arcana:
                        continue
                    if not m and p1.pid > p2.pid:
                        continue
                    if self._check_comp2(t, p1, p2):
                        result.append((t, p1, p2))
        else:
            # 通过素材计算合成目标
            pass
        self.result = result

    @staticmethod
    def _check_comp3(target, p1, p2, p3):
        """
        三体合成算法测试结果：3个素材p的等级总和<(目标p初始等级-4)*3即成立
        :param target: 目标p
        :param p1: 素材1
        :param p2: 素材2
        :param p3: 素材3
        :return: 符合三体合成返回True，否则返回False
        """
        if (target.pid != p1.pid != p2.pid != p3.pid and p1.pid != p3.pid != target.pid != p2.pid
                and target.special == 0 and p1.level + p2.level + p3.level < (target.level - 4) * 3):
            for p in persona:
                if (target.pid > p.pid != p3.pid and p1.pid != p.pid != p2.pid
                        and p.special == 0 and p.arcana == target.arcana):
                    if p1.level + p2.level + p3.level < (p.level - 4) * 3:
                        return False
            else:
                return True

    def calculate_persona_by_comp3_arithmetic(self):
        """
        通过三体合成算法计算人格面具
        :return:
        """
        t = self.target
        m1 = self.material1
        m2 = self.material2
        m3 = self.material3
        result = []
        if not (t or m1 or m2 or m3):
            self.result = result
            return
        if m3 and not m2:
            m2, m3 = m3, m2
        if t:
            # 计算合成素材
            if (t.special == 1 or (m1 and m1.pid == t.pid) or (m2 and m2.pid == t.pid)
                    or (m3 and m3.pid == t.pid) or (m1 and m2 and m1.pid == m2.pid)
                    or (m2 and m3 and m2.pid == m3.pid)):
                self.result = result
                return
            """用于记录三体合成的arcana公式的四元组, 0:结果, 1:p1, 2:p2, 3:p3"""
            step1 = {}
            """第一步中间过程，键是可用的p1的arcana，值是对应的p2和p3的arcana的二元组列表"""
            for c3 in comp3:
                if c3['result'] == t.arcana:
                    # 如果设置了素材1，则step1的键只有素材1的arcana
                    if m1 and m1.arcana != c3['arcana1']:
                        continue
                    if m1 and m1.arcana == c3['arcana1']:
                        if not step1.get(m1.arcana):
                            step1[m1.arcana] = []
                        for c2 in comp2:
                            if c2['result'] == c3['arcana2']:
                                step1[m1.arcana].append((c2['arcana1'], c2['arcana2']))
                    else:
                        for arcana1 in Arcana:
                            if c3['result'] == t.arcana and c3['arcana1'] == arcana1:
                                if not step1.get(arcana1):
                                    step1[arcana1] = []
                                for c2 in comp2:
                                    if c2['result'] == c3['arcana2']:
                                        step1[arcana1].append((c2['arcana1'], c2['arcana2']))
            arcana_formula = []
            set_arcana1 = set(step1.keys())
            """用于保存允许的p1的arcana类型"""
            for a1 in step1:
                for a2 in step1[a1]:
                    arcana_formula.append((t.arcana, a1, a2[0], a2[1]))
            for p1 in persona:
                if (m1 and p1 != m1) or (p1.arcana not in set_arcana1):
                    continue
                set_arcana2 = set()
                """用于保存当前p1允许的p2的arcana类型"""
                if m2:
                    set_arcana2.add(m2.arcana)
                else:
                    for row in arcana_formula:
                        if row[1] == p1.arcana and row[2] <= row[3]:
                            set_arcana2.add(row[2])
                for p2 in persona:
                    if (m2 and p2 != m2) or (p2.arcana not in set_arcana2):
                        continue
                    set_arcana3 = set()
                    # 根据p1和p2计算p3的arcana
                    for row in arcana_formula:
                        if row[1] == p1.arcana and row[2] == p2.arcana:
                            if m2:
                                set_arcana3.add(row[3])
                            else:
                                if row[2] <= row[3]:
                                    set_arcana3.add(row[3])
                    for p3 in persona:
                        if (m3 and p3 != m3) or (p3.arcana not in set_arcana3):
                            continue
                        if not m2 and p2.arcana == p3.arcana and p2.pid > p3.pid:
                            continue
                        if Calc._check_comp3(t, p1, p2, p3):
                            result.append((t, p1, p2, p3))
        else:
            # 计算合成结果
            pass
        self.result = result

    @classmethod
    def find_persona_by_name(cls, name):
        print(''.center(20, '-'))
        for p in persona:
            if name in p.name:
                Calc.showPersona(p)
                print(''.center(20, '-'))
        input('(回车确认)')

    @classmethod
    def find_persona_by_arcana(cls, arcana):
        if not arcana:
            return
        arcana = Arcana.parse(arcana)
        print(''.center(20, '-'))
        for p in persona:
            if p.arcana == arcana:
                Calc.showPersona(p)
                print(''.center(20, '-'))
        input('(回车确认)')

    @classmethod
    def find_persona_by_level(cls, low, high):
        print(''.center(20, '-'))
        for p in persona:
            if low <= p.level <= high:
                Calc.showPersona(p)
                print(''.center(20, '-'))
        input('(回车确认)')

    @classmethod
    def find_persona_by_position(cls, position):
        if position == 1:
            position = '城堡'
        elif position == 2:
            position = '浴场'
        elif position == 3:
            position = '剧场'
        elif position == 4:
            position = '迷宫'
        elif position == 5:
            position = '研究所'
        elif position == 6:
            position = '乐土'
        elif position == 7:
            position = '稻羽'
        elif position == 8:
            position = '空洞'
        elif position == 9:
            position = '黄泉'
        print(''.center(20, '-'))
        for p in persona:
            if position in p.position:
                Calc.showPersona(p)
                print(''.center(20, '-'))
        input('(回车确认)')

    def find_formula_by_comp2(self):
        """
        查询二体合成公式
        以self.target为目标计算合成公式
        如果self.material1是None，计算所有合成公式
        如果self.material1不为None，计算以self.material1为素材的合成公式
        :return:返回以target为目标，material1为素材的合成公式
        """
        if not self.target:
            raise IndexError('未设置合成目标')
        if self.target.special == 1:
            raise IndexError('特殊合成面具不可作为二体/三体合成目标')
        self.calculate_persona_by_comp2_arithmetic()
        self.show_result()

    def find_formula_by_comp3(self):
        if not self.target:
            raise IndexError('未设置合成目标')
        if self.target.special == 1:
            raise IndexError('特殊合成面具不可作为二体/三体合成目标')
        self.calculate_persona_by_comp3_arithmetic()
        self.show_result()


def mainMenu():
    print('主菜单'.center(20, '='))
    print('1.查询人格面具')
    print('2.查询合成公式')
    print('0.退出')
    print('请选择'.center(20, '-'))


def pMenu():
    """查询人格面具菜单"""
    while True:
        print('查询人格面具'.center(20, '='))
        print('1.按名称查询')
        print('2.按阿尔卡那查询')
        print('3.按等级查询')
        print('4.按掉落位置查询')
        print('0.返回')
        print('请选择'.center(20, '-'))
        try:
            c = int(input())
            if c == 0:
                break
            elif c == 1:
                name = input('请输入查询的名称:\n')
                Calc.find_persona_by_name(name)
            elif c == 2:
                arcana = input('请输入查询的阿尔卡那:\n')
                Calc.find_persona_by_arcana(arcana)
            elif c == 3:
                low = int(input('请输入查询的最低等级:'))
                high = int(input('请输入查询的最高等级:'))
                if low > high:
                    input('最低等级应小于等于最高等级(回车确认)')
                elif low < 0 or high < 0:
                    input('等级应大于0')
                else:
                    Calc.find_persona_by_level(low, high)
            elif c == 4:
                print('请选择查询的位置:')
                print('1.城堡')
                print('2.浴场')
                print('3.剧场')
                print('4.迷宫')
                print('5.研究所')
                print('6.乐土')
                print('7.稻羽')
                print('8.空洞')
                print('9.黄泉')
                print('请选择'.center(20, '-'))
                position = int(input())
                if position in (1, 2, 3, 4, 5, 6, 7, 8, 9):
                    Calc.find_persona_by_position(position)
            else:
                input('未定义的选项(回车确认)')
        except ValueError:
            input('请输入数字(回车确认)')
        except TypeError as e:
            print(e, end='')
            input('(回车确认)')


def compMenu(calc):
    """查询合成菜单"""
    while True:
        print('查询合成公式'.center(20, '='))
        print('当前合成目标\t:', calc.target)
        print('当前素材1\t:', calc.material1)
        print('当前素材2\t:', calc.material2)
        print('当前素材3\t:', calc.material3)
        print(''.center(20, '-'))
        print('1.设置合成目标')
        print('2.二体合成')
        print('3.三体合成')
        print('4.设置素材1')
        print('5.设置素材2')
        print('6.设置素材3')
        print('7.清除素材1')
        print('8.清除素材2')
        print('9.清除素材3')
        print('0.返回')
        print('请选择'.center(20, '-'))
        try:
            c = input()
            if c == '0':
                break
            elif c == '1':
                print('请输入合成目标(留空不做改变):')
                name = input()
                calc.target = name
            elif c == '2':
                if not calc.target:
                    input('必须设置合成目标(回车确认)')
                    continue
                calc.find_formula_by_comp2()
            elif c == '3':
                if not calc.target:
                    input('必须设置合成目标(回车确认)')
                    continue
                calc.find_formula_by_comp3()
            elif c == '4':
                print('请输入素材1(留空不做改变):')
                name = input()
                calc.material1 = name
            elif c == '5':
                print('请输入素材2(留空不做改变):')
                name = input()
                calc.material2 = name
            elif c == '6':
                print('请输入素材3(留空不做改变):')
                name = input()
                calc.material3 = name
            elif c == '7':
                del calc.material1
            elif c == '8':
                del calc.material2
            elif c == '9':
                del calc.material3
            else:
                input('未定义的选项(回车确认)')
        except Exception as e:
            print(e, end='')
            input('(回车确认)')


def main():
    init()
    calc = Calc()
    while True:
        try:
            mainMenu()
            choice = int(input())
            if choice == 0:
                break
            elif choice == 1:
                pMenu()
            elif choice == 2:
                compMenu(calc)
            else:
                input('未定义的选项(回车确认)')
        except ValueError:
            input('请输入数字(回车确认)')


if __name__ == '__main__':
    main()
